import torch
import torchvision.datasets
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

input = torch.tensor([
    [1, 2, 0, 3, 1],
    [0, 1, 2, 3, 1],
    [1, 2, 1, 0, 0],
    [5, 2, 3, 1, 1],
    [2, 1, 0, 1, 1]
], dtype=torch.float32)

input = torch.reshape(input, (-1, 1, 5, 5))
"""
最大池化，相当于是将图片压缩了， 在保留图片特征的情况下减少了像素
"""
output = torch.nn.MaxPool2d(kernel_size=3, ceil_mode=True)(input)
print(output)

class AModule(torch.nn.Module):
    def __init__(self, cell_mode=False):
        super(AModule, self).__init__()
        self.pool = torch.nn.MaxPool2d(kernel_size=3, ceil_mode=cell_mode)

    def forward(self, input):
        return self.pool(input)

# cell_mode 为 true， 可以认为，相对于 false 来说，压缩的没那么厉害
module = AModule(cell_mode=True)
false_module = AModule(cell_mode=False)

output2 = false_module(input)
print("-----")
print(output2)


print("====")

dataset = torchvision.datasets.CIFAR10(root="datasets", train=False, transform=torchvision.transforms.ToTensor(), download=True)
dataloader = DataLoader(dataset=dataset, batch_size=64, shuffle=True)

step = 0
writer = SummaryWriter("logs")
for imgs, targets in dataloader:
    output1 = module(imgs)
    output2 = false_module(imgs)
    writer.add_images("original", imgs, step)
    writer.add_images("cell_true", output1, step)
    writer.add_images("cell_false", output2, step)
    step = step + 1

writer.close()
print("success")